package com.kingsoft.minibase;

import java.io.Closeable;
import java.io.IOException;
import java.util.Iterator;
import java.util.SortedMap;
import java.util.concurrent.ConcurrentSkipListMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.locks.ReentrantReadWriteLock;

import com.kingsoft.minibase.MiniBase.Flusher;
import com.kingsoft.minibase.MiniBase.Compactor;
import com.kingsoft.minibase.MStore.SeekIter;

import org.apache.log4j.Logger;

/**
 * 负责读写数据、刷盘
 */
public class MemStore implements Closeable {
    private static final Logger LOG = Logger.getLogger(MemStore.class);
    private Config config;
    private AtomicLong dataSize = new AtomicLong();

    private volatile ConcurrentSkipListMap<KeyValue, KeyValue> kvMap;
    private volatile ConcurrentSkipListMap<KeyValue, KeyValue> snapshot;
    private AtomicBoolean isSnapshotting = new AtomicBoolean(false);
    private ReentrantReadWriteLock updateLock = new ReentrantReadWriteLock();
    private ExecutorService pool;
    private Flusher flusher;
    private Compactor compactor;

    public MemStore(Config config, Flusher flusher, ExecutorService pool) {
        this.config = config;
        this.flusher = flusher;
        this.pool = pool;

        dataSize.set(0);
        kvMap = new ConcurrentSkipListMap<>();
        snapshot = null;
    }

    // 写入数据，将keyValue对象添加到内存中，同时更改内存中的数据大小dataSize
    public void add(KeyValue keyValue) throws IOException {
        // 因为多线程并发写入的情况，在进行add操作的前后都需要一次刷写判断
        flushIfNeeded(true);
        // 加锁
        updateLock.writeLock().lock();
        // 判断是新插入的key还是更改key的值
        try {
            KeyValue preKeyValue;
            if ((preKeyValue = kvMap.put(keyValue, keyValue)) == null) {
                dataSize.addAndGet(keyValue.getSerializeSize());
            } else {
                dataSize.addAndGet(keyValue.getSerializeSize() - preKeyValue.getSerializeSize());
            }
        } finally {
            updateLock.writeLock().unlock();
        }
        flushIfNeeded(false);
    }

    // 需要查看是否正在进行flush操作
    public void flushIfNeeded(boolean shouldBlocking) throws IOException {
        if (getDataSize() > config.getMaxMemstoreSize()) {
            if (isSnapshotting.get() && shouldBlocking) {
                throw new IOException(
                        "Memstore is full, currentDataSize=" + dataSize.get() + "B, maxMemstoreSize="
                                + config.getMaxMemstoreSize() + "B, please wait until the flushing is finished.");
            } else if (isSnapshotting.compareAndSet(false, true)){
                pool.submit(new FlusherTask());
            }
        }
    }

    public long getDataSize() {
        return dataSize.get();
    }

    private class FlusherTask implements Runnable {
        @Override
        public void run() {
            // 1. 创建snapshot，加锁
            updateLock.writeLock().lock();
            try {
                snapshot = kvMap;
                kvMap = new ConcurrentSkipListMap<>();
                dataSize.set(0);
            } finally {
                updateLock.writeLock().unlock();
            }

            // 2. 刷写到DiskFile
            boolean success = true;
            for (int i = 0; i < config.getFlushMaxRetries(); i++) {
                try {
                    // 把跳表包装成迭代器
                    flusher.flush(new IteratorWrapper(snapshot));
                } catch (IOException e) {
                    success = false;
                    LOG.error("Failed to flush memstore, retries=" + i + ", maxFlushRetries="
                                    + config.getFlushMaxRetries(),
                            e);
                    if (i > config.getFlushMaxRetries()) {
                        break;
                    }
                }
            }

            // 3. 清理snapshot
            if (success) {
                snapshot = null;
                isSnapshotting.compareAndSet(true, false);
            }
        }
    }

    public static class IteratorWrapper implements SeekIter<KeyValue> {
        private SortedMap<KeyValue, KeyValue> sortedMap;
        private Iterator<KeyValue> it;

        public IteratorWrapper(SortedMap sortedMap) {
            this.sortedMap = sortedMap;
            this.it = sortedMap.values().iterator();
        }
        @Override
        public boolean hasNest() throws IOException {
            return it != null && it.hasNext();
        }

        @Override
        public KeyValue next() throws IOException {
            return it.next();
        }

        @Override
        public void seekTo(KeyValue kv) {
            // 更新it，只保留比kv小的值
            it = sortedMap.tailMap(kv).values().iterator();
        }
    }


    @Override
    public void close() throws IOException {

    }
}
